#include "ip_packet.h"

#include <stdint.h>
#include <unordered_set>
#include "RawPacket.h"
#include "Packet.h"
#include "ip_reassembly.h"
#include "IPv4Layer.h"
#include "EthLayer.h"
#include "VlanLayer.h"
#include "TcpLayer.h"
#include "UdpLayer.h"
#include "server/server.h"

namespace secsmart_tcpip
{
IpPacketHandler::IpPacketHandler()
{
    init();
}

IpPacketHandler::~IpPacketHandler()
{
    delete m_re_pkt_;
    m_re_pkt_ = nullptr;

    delete m_frag_handler_;
    m_frag_handler_ = nullptr;
}

void IpPacketHandler::init()
{
#define IP_FRAG_STREAM_DEFAULT_MAX_SIZE 10000 // ip 分片流缺省规格
#define IP_FRAG_CACHE_DEFAULT_MAX_SIZE 30000 // ip 分片缓存缺省规格
#define IP_FRAG_DEFAULT_TIMEOUT 3000 // 分片缺省老化时间3秒

    // 报文过滤初始化
    m_enable_filter_packet = false;

    // 分片处理初始化
    m_frag_handler_ = new IPReassembly;
    if (m_frag_handler_ != nullptr) {
        m_enable_frag_reassembly_ = true;
        m_frag_handler_->set_cfg(IP_FRAG_STREAM_DEFAULT_MAX_SIZE, IP_FRAG_CACHE_DEFAULT_MAX_SIZE,
                                 IP_FRAG_DEFAULT_TIMEOUT, false);
        m_re_pkt_ = nullptr;
    }
}

void IpPacketHandler::set_frag_reassembly_enable(bool on_off)
{
    m_enable_frag_reassembly_ = on_off;
}

void IpPacketHandler::set_frag_reassembly_cfg(uint32_t ip_stream_size, uint32_t frag_cache_size,
    uint32_t timeout, bool detect_ooo)
{
    if (m_frag_handler_ != nullptr) {
        m_frag_handler_->set_cfg(ip_stream_size, frag_cache_size, timeout, detect_ooo);
    }
}

void IpPacketHandler::get_frag_reassembly_cfg(uint32_t &ip_stream_size, uint32_t &frag_cache_size,
    uint32_t &timeout, bool &detect_ooo)
{
    if (m_frag_handler_ == nullptr) {
        return;
    }

    IPReassembly::IPReassemblyCfg *cfg = m_frag_handler_->get_cfg();
    if (cfg == nullptr) {
        return;
    }
    ip_stream_size = cfg->max_stream_num;
    frag_cache_size = cfg->max_frag_cache_num;
    timeout = cfg->timeout;
    detect_ooo = cfg->detect_ooo_enable;
}

uint32_t IpPacketHandler::process_packet(uint8_t *in_pkt, uint8_t **out_pkt, packet_meta_t *meta)
{
    if (in_pkt == NULL || out_pkt == NULL || meta == NULL) {
        return IP_PKT_ERR;
    }
    *out_pkt = NULL;
    m_stat.stat_rcv_pkt(meta->pkt_len);

    // 初始化报文
    RawPacket raw_pkt(in_pkt, meta->pkt_len, meta->timestamp, false, meta->link_type);
    Packet pkt;
    pkt.init(&raw_pkt, false, meta);

    // 进行报文解析
    pkt.parsePacket(TCP | UDP);

    if (meta->ip_len + meta->ip_offset > meta->pkt_len) {
        m_stat.stat_drop_pkt(IP_PKT_MALFORMED);
        return IP_PKT_MALFORMED;
    }

    // 分片报文处理
    if(pkt.isFragmentPacket()) {
        uint32_t ret = process_frag_packet(&pkt);
        if (ret != IP_PKT_OK) {
            m_stat.stat_drop_pkt(ret);
            return ret;
        }
        *out_pkt = m_re_pkt_->getRawPacketData();
    } else {
        *out_pkt = in_pkt;
    }

    // 非tcp udp报文，直接返回
    if ((meta->ip_protocol != PACKETPP_IPPROTO_UDP) && 
        (meta->ip_protocol != PACKETPP_IPPROTO_TCP)) {
        m_stat.stat_drop_pkt(IP_PKT_UNSUPPORT);
        return IP_PKT_UNSUPPORT;
    }

    // 进行报文过滤
    if (m_enable_filter_packet && filter_packet(meta)) {
        m_stat.stat_drop_pkt(IP_PKT_FILTER_DROP);
        return  IP_PKT_FILTER_DROP;
    }
    
    // 报文解析处理完成
    m_stat.stat_snd_pkt(meta->pkt_len);
    return IP_PKT_OK;
}

/**
 * @brief 快速解析报文，仅支持eth[vlan] + ipv4 + tcp/udp 的报文
 * @param pkt 输入的报文
 * @param meta 报文分析结果
 * @return true 解析OK，false 需要进一步处理
 */
bool IpPacketHandler::fast_parse_packet(uint8_t *pkt, packet_meta_t *meta)
{
    if (meta->link_type != LINKTYPE_ETHERNET) {
        return false;
    }
    //!TODO: 报文合法性验证
    ether_header *eth_hdr = (ether_header *)pkt;
    vlan_header *vlan_hdr = NULL;
    iphdr *ip_hdr = NULL;

    if (eth_hdr->etherType == PCPP_ETHERTYPE_IP) { // ipv4 报文
        ip_hdr = (iphdr*)(eth_hdr + 1);
    } else if (eth_hdr->etherType == PCPP_ETHERTYPE_VLAN) { // 快速解析最多支持一层vlan
        vlan_hdr = (vlan_header*)(eth_hdr + 1);
        if (vlan_hdr->etherType != PCPP_ETHERTYPE_IP) {
            return false;
        }
        ip_hdr = (iphdr*)(vlan_hdr + 1);
    } else {
        return false;
    }

    uint32_t ip_hdr_len = ip_hdr->internetHeaderLength * 4;
    tcphdr * tcp_hdr = NULL;
    udphdr *udp_hdr = NULL;

    if (ip_hdr->protocol == PACKETPP_IPPROTO_TCP) { // tcp 报文
        tcp_hdr = (tcphdr*)((uint8_t*)ip_hdr + ip_hdr_len);
    } else if (ip_hdr->protocol == PACKETPP_IPPROTO_UDP) { // udp 报文
        udp_hdr = (udphdr*)((uint8_t*)ip_hdr + ip_hdr_len);
        if (be16toh(udp_hdr->portDst) == 4789) { // vxlan 报文
            return false;
        }
    } else {
        return false;
    }

    if (ip_hdr->fragmentOffset & (uint16_t)0xFF3F) { // 分片报文
        return false;
    }

    // 收集信息
    meta->smac = MacAddress(eth_hdr->srcMac);
    meta->dmac = MacAddress(eth_hdr->dstMac);
    if (vlan_hdr != NULL) {
        meta->vlanid = be16toh(vlan_hdr->vlan);
    }
    meta->ip_version = 4;
    meta->sip = IPv4Address(ip_hdr->ipSrc);
    meta->dip = IPv4Address(ip_hdr->ipDst);
    meta->ip_protocol = ip_hdr->protocol;
    if (tcp_hdr != NULL) {
        meta->sport = be16toh(tcp_hdr->portSrc);
        meta->dport = be16toh(tcp_hdr->portDst);
    } else {
        meta->sport = be16toh(udp_hdr->portSrc);
        meta->dport = be16toh(udp_hdr->portDst);
    }
    return true;
}

uint32_t IpPacketHandler::process_frag_packet(Packet *pkt)
{
    if (m_enable_frag_reassembly_ == false || m_frag_handler_ == nullptr) { //分片重组没有开启，直接丢弃
        return IP_PKT_FRAG_DROP;
    }

    uint32_t status;
    Packet*new_pkt = m_frag_handler_->process_packet(pkt, status, TCP | UDP);
    if (new_pkt == nullptr) { // 分片报文还不完整
        if (status == IPReassembly::REASSEMBLY_CACHE_FRAG) {
            return IP_PKT_FRAG_CACHED;
        } else {
            return IP_PKT_FRAG_DROP;
        }
    }
    delete m_re_pkt_;
    m_re_pkt_ = new_pkt;

    // 重组后的报文还是分片报文，需要继续重组.外层是隧道报文，内层是IP分片报文的场景
    if (m_re_pkt_->isFragmentPacket()) { 
        new_pkt = m_frag_handler_->process_packet(m_re_pkt_, status, TCP | UDP);
        if (new_pkt == nullptr) { // 分片报文还不完整
            if (status == IPReassembly::REASSEMBLY_CACHE_FRAG) {
                return IP_PKT_FRAG_CACHED;
            } else {
                return IP_PKT_FRAG_DROP;
            }
        }
        delete m_re_pkt_;
        m_re_pkt_ = new_pkt;
    }
    return IP_PKT_OK;
}

// 返回true，表示报文被过滤掉
bool IpPacketHandler::filter_packet(packet_meta_t *meta)
{
    // 检查报文的源目(IP port)对是否为server的监听地址，是监听地址，返回false，否则返回true
    if (ServerManager::get_instance().is_work_server(meta->sip, meta->sport, meta->ip_protocol)) {
        meta->hit_server = 1;
        meta->direction = PKT_DIRECTION_SERVER_TO_CLIENT;
        return false;
    }

    if (ServerManager::get_instance().is_work_server(meta->dip, meta->dport, meta->ip_protocol)) {
        meta->hit_server = 1;
        meta->direction = PKT_DIRECTION_CLIENT_TO_SERVER;
        return false;
    }

    return true;
}

}

